{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RL 训练v4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
      "Unsloth: Failed to patch Gemma3ForConditionalGeneration.\n",
      "🦥 Unsloth Zoo will now patch everything to make training faster!\n",
      "INFO 04-22 21:14:37 [__init__.py:239] Automatically detected platform cuda.\n",
      "==((====))==  Unsloth 2025.3.19: Fast Qwen2 patching. Transformers: 4.51.3. vLLM: 0.8.4.\n",
      "   \\\\   /|    NVIDIA GeForce RTX 4090. Num GPUs = 1. Max memory: 22.159 GB. Platform: Linux.\n",
      "O^O/ \\_/ \\    Torch: 2.6.0+cu124. CUDA: 8.9. CUDA Toolkit: 12.4. Triton: 3.2.0\n",
      "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.29.post2. FA2 = False]\n",
      " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
      "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n",
      "Unsloth: vLLM loading /data/countdown/output/models/Qwen2.5-1.5B-Instruct with actual GPU utilization = 39.23%\n",
      "Unsloth: Your GPU has CUDA compute capability 8.9 with VRAM = 22.16 GB.\n",
      "Unsloth: Using conservativeness = 1.0. Chunked prefill tokens = 4096. Num Sequences = 192.\n",
      "Unsloth: vLLM's KV Cache can use up to 5.68 GB. Also swap space = 6 GB.\n",
      "INFO 04-22 21:14:47 [config.py:689] This model supports multiple tasks: {'score', 'reward', 'embed', 'generate', 'classify'}. Defaulting to 'generate'.\n",
      "Unsloth: vLLM Bitsandbytes config using kwargs = {'load_in_8bit': False, 'load_in_4bit': True, 'bnb_4bit_compute_dtype': 'bfloat16', 'bnb_4bit_quant_storage': 'uint8', 'bnb_4bit_quant_type': 'fp4', 'bnb_4bit_use_double_quant': False, 'llm_int8_enable_fp32_cpu_offload': False, 'llm_int8_has_fp16_weight': False, 'llm_int8_skip_modules': [], 'llm_int8_threshold': 6.0}\n",
      "INFO 04-22 21:14:48 [llm_engine.py:243] Initializing a V0 LLM engine (v0.8.4) with config: model='/data/countdown/output/models/Qwen2.5-1.5B-Instruct', speculative_config=None, tokenizer='/data/countdown/output/models/Qwen2.5-1.5B-Instruct', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=4096, download_dir=None, load_format=bitsandbytes, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=bitsandbytes, enforce_eager=False, kv_cache_dtype=auto,  device_config=cuda:0, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=/data/countdown/output/models/Qwen2.5-1.5B-Instruct, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=True, chunked_prefill_enabled=False, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={\"level\":0,\"splitting_ops\":[],\"compile_sizes\":[],\"cudagraph_capture_sizes\":[192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],\"max_capture_size\":192}, use_cached_outputs=False, \n",
      "INFO 04-22 21:14:48 [cuda.py:292] Using Flash Attention backend.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[rank0]:[W422 21:14:57.287021590 ProcessGroupGloo.cpp:727] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-22 21:14:57 [parallel_state.py:959] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0\n",
      "INFO 04-22 21:14:57 [model_runner.py:1110] Starting to load model /data/countdown/output/models/Qwen2.5-1.5B-Instruct...\n",
      "INFO 04-22 21:14:58 [loader.py:1166] Loading weights with BitsAndBytes quantization. May take a while ...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38080299c3444bd2abbc0c1193a12f4e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-22 21:14:58 [punica_selector.py:18] Using PunicaWrapperGPU.\n",
      "INFO 04-22 21:14:59 [model_runner.py:1146] Model loading took 1.2132 GiB and 0.954067 seconds\n",
      "INFO 04-22 21:15:01 [worker.py:267] Memory profiling takes 1.72 seconds\n",
      "INFO 04-22 21:15:01 [worker.py:267] the current vLLM instance can use total_gpu_memory (22.16GiB) x gpu_memory_utilization (0.39) = 8.69GiB\n",
      "INFO 04-22 21:15:01 [worker.py:267] model weights take 1.21GiB; non_torch_memory takes 0.08GiB; PyTorch activation peak memory takes 1.05GiB; the rest of the memory reserved for KV Cache is 6.35GiB.\n",
      "INFO 04-22 21:15:01 [executor_base.py:112] # cuda blocks: 14861, # CPU blocks: 14043\n",
      "INFO 04-22 21:15:01 [executor_base.py:117] Maximum concurrency for 4096 tokens per request: 58.05x\n",
      "INFO 04-22 21:15:04 [model_runner.py:1456] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2bb9853304d640ce8077508fadfe3c56",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Capturing CUDA graph shapes:   0%|          | 0/27 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 04-22 21:15:34 [model_runner.py:1598] Graph capturing finished in 30 secs, took 2.87 GiB\n",
      "INFO 04-22 21:15:34 [llm_engine.py:449] init engine (profile, create kv cache, warmup model) took 35.01 seconds\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sliding Window Attention is enabled but not implemented for `eager`; unexpected results may be encountered.\n",
      "Unsloth 2025.3.19 patched 28 layers with 28 QKV layers, 28 O layers and 28 MLP layers.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "# fix：https://github.com/unslothai/unsloth/issues/2299#issuecomment-2782067709\n",
    "os.environ[\"VLLM_USE_V1\"] = '0'\n",
    "# 国内需要禁止统计，否则会卡在模型加载的地方（连不到外网）\n",
    "os.environ[\"UNSLOTH_DISABLE_STATISTICS\"] = \"0\"\n",
    "\n",
    "from unsloth import FastLanguageModel, is_bfloat16_supported\n",
    "import torch\n",
    "\n",
    "max_seq_length = 4096 # Can increase for longer reasoning traces\n",
    "lora_rank = 64 # Larger rank = smarter, but slower\n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name = \"/data/countdown/output/models/Qwen2.5-1.5B-Instruct\", # change to your model path\n",
    "    # 使用 merge.ipynb 合并后的模型作为 RL 的基座，从而解决冷启动的问题。\n",
    "    # model_name = \"/home/cynic/models/qwen2.5-1.5b-sft-distill-merged\",\n",
    "    max_seq_length = max_seq_length,\n",
    "    load_in_4bit = True, # False for LoRA 16bit\n",
    "    local_files_only=True,\n",
    "    fast_inference = True, # Enable vLLM fast inference\n",
    "    max_lora_rank = lora_rank,\n",
    "    gpu_memory_utilization = 0.4, # Reduce if out of memory\n",
    ")\n",
    "\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128\n",
    "    target_modules = [\n",
    "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "        \"gate_proj\", \"up_proj\", \"down_proj\",\n",
    "    ], # Remove QKVO if out of memory\n",
    "    lora_alpha = lora_rank * 2,\n",
    "    use_gradient_checkpointing = \"unsloth\", # Enable long context finetuning\n",
    "    random_state = 3407,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'numbers': [91, 100, 44],\n",
       " 'target': 147,\n",
       " 'ground_truth_solution': '(91 - 44) + 100',\n",
       " 'prompt': [{'content': 'You are a helpful assistant. You first thinks about the reasoning process in the mind and then provides the user with the answer.',\n",
       "   'role': 'system'},\n",
       "  {'content': 'Using the numbers 91, 100, 44, create an equation that equals 147. You can use basic arithmetic operations (+, -, *, /) one or multiple times but each number can only be used once, and you must use all the numbers. Show your work in <think> </think> tags. And return the final equation in <answer> </answer> tags, for example <answer>(1 + 2) / 3</answer>. Think step by step inside <think> tags.',\n",
       "   'role': 'user'}]}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from constant import SYSTEM_PROMPT, USER_PROMPT_TPL, parse_user_prompt\n",
    "from datasets import load_dataset\n",
    "def get_countdown_questions(data_file):\n",
    "    data = load_dataset(\"json\", data_files=data_file)[\"train\"]\n",
    "    data = data.map(lambda x: {\n",
    "        'prompt': [\n",
    "            {'role': 'system', 'content': SYSTEM_PROMPT},\n",
    "            {'role': 'user', \"content\": parse_user_prompt(USER_PROMPT_TPL, x[\"numbers\"], x[\"target\"])}\n",
    "        ],\n",
    "    })\n",
    "    return data\n",
    "\n",
    "dataset = get_countdown_questions(\"data/rl_data_simple_10k.jsonl\")\n",
    "\n",
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "def extract_xml_answer(response):\n",
    "    pattern = r\"<answer>(.*?)</answer>\"\n",
    "    match = re.search(pattern, response, re.DOTALL)\n",
    "    return match.group(1) if match else \"\"\n",
    "\n",
    "# Reward function that checks if the completion follows a more relaxed format\n",
    "def soft_format_reward_func(completions, **kwargs) -> list[float]:\n",
    "    pattern = r\"<think>.*?</think>\\s*<answer>.*?</answer>\"\n",
    "    responses = [completion[0][\"content\"] for completion in completions]\n",
    "    matches = [re.match(pattern, r, re.DOTALL) for r in responses]\n",
    "    return [0.5 if match else 0.0 for match in matches]\n",
    "\n",
    "# 正确性奖励\n",
    "def correctness_reward_func(completions, numbers, target, **kwargs) -> list[float]:\n",
    "    responses = [completion[0]['content'] for completion in completions]\n",
    "    extracted_responses = [extract_xml_answer(r) for r in responses]\n",
    "\n",
    "    def correct(numbers, target, solution):\n",
    "        try:\n",
    "            if '=' in solution:\n",
    "                # 删除掉等号以及等号后的内容\n",
    "                solution = solution.split('=')[0]\n",
    "            if not eval(solution) == target:\n",
    "                # 如果 solution 计算结果不等于 target，则认为不正确\n",
    "                return False, \"Solution is not equal to target\"\n",
    "            # Define a regex pattern that only allows numbers, operators, parentheses, and whitespace\n",
    "            allowed_pattern = r'^[\\d+\\-*/().\\s]+$'\n",
    "            if not re.match(allowed_pattern, solution):\n",
    "                # 如果 solution 不符合要求，则认为不正确\n",
    "                return False, \"Solution is not allowed\"\n",
    "            # 检查 solution 是否使用了所有数字\n",
    "            used_numbers = [int(n) for n in re.findall(r'\\d+', solution)]\n",
    "            if sorted(used_numbers) != sorted(numbers):\n",
    "                # 如果 solution 没有使用所有数字没有仅使用一次，则认为不正确\n",
    "                return False, \"Solution is not using all numbers or each number is not used only once\"\n",
    "            return True, \"Solution is correct\"\n",
    "        except Exception as e:\n",
    "            # 计算失败，则认为不正确\n",
    "            return False, f\"Solution calculation failed: {e}\"\n",
    "    \n",
    "    scores = []\n",
    "    for i, r in enumerate(extracted_responses):\n",
    "        scores.append(2 if correct(numbers[i], target[i], r)[0] else 0)\n",
    "    \n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'target': 656,\n",
       " 'numbers': [85, 36, 16, 519],\n",
       " 'prompt': 'Using the numbers 85, 36, 16, 519, create an equation that equals 656. You can use basic arithmetic operations (+, -, *, /) one or multiple times but each number can only be used once, and you must use all the numbers. Show your work in <think> </think> tags. And return the final equation in <answer> </answer> tags, for example <answer>(1 + 2) / 3</answer>. Think step by step inside <think> tags.',\n",
       " 'completion': \"<think>Okay, let's see. I need to make an equation using 85, 36, 16, and 519, each exactly once, with basic operations, and the result should be 656. Hmm.\\n\\nFirst, maybe I should look at the target number, 656. Let me think if there's any way to break it down into parts that could relate to these numbers. For instance, 656 divided by something... Maybe 656 is close to 519 plus something? Wait, 519 is already given as part of the numbers. So maybe 519 plus what else would get me to 656?\\n\\nWait, 519 + 137 = 656. But how do I get 137 from the remaining numbers: 85, 36, 16. Let's check. The largest number here is 85. If I subtract 85 from 137, that gives 52. Not sure. Alternatively, adding 85 and 36 gives 121, which is still less than 137. Then 121 + 16 is 137. Oh! That works. So 85 + 36 + 16 = 137. Then 519 + 137 = 656. But wait, does that use all four numbers? Yes: 519, 85, 36, 16. Each used once. So the equation would be 519 + (85 + 36 + 16) = 656. Let me verify:\\n\\n85 + 36 = 121\\n\\n121 + 16 = 137\\n\\n519 + 137 = 656. Perfect!\\n\\nSo that seems to work. Let me just double-check if there's another possible combination, but this seems straightforward. Alternatively, maybe using multiplication or division somewhere, but perhaps not necessary here.\\n</think>\\n\\n<answer>519 + 85 + 36 + 16 = 656</answer>\",\n",
       " 'reasoning': \"Okay, let's see. I need to make an equation using 85, 36, 16, and 519, each exactly once, with basic operations, and the result should be 656. Hmm.\\n\\nFirst, maybe I should look at the target number, 656. Let me think if there's any way to break it down into parts that could relate to these numbers. For instance, 656 divided by something... Maybe 656 is close to 519 plus something? Wait, 519 is already given as part of the numbers. So maybe 519 plus what else would get me to 656?\\n\\nWait, 519 + 137 = 656. But how do I get 137 from the remaining numbers: 85, 36, 16. Let's check. The largest number here is 85. If I subtract 85 from 137, that gives 52. Not sure. Alternatively, adding 85 and 36 gives 121, which is still less than 137. Then 121 + 16 is 137. Oh! That works. So 85 + 36 + 16 = 137. Then 519 + 137 = 656. But wait, does that use all four numbers? Yes: 519, 85, 36, 16. Each used once. So the equation would be 519 + (85 + 36 + 16) = 656. Let me verify:\\n\\n85 + 36 = 121\\n\\n121 + 16 = 137\\n\\n519 + 137 = 656. Perfect!\\n\\nSo that seems to work. Let me just double-check if there's another possible combination, but this seems straightforward. Alternatively, maybe using multiplication or division somewhere, but perhaps not necessary here.\\n\",\n",
       " 'solution': '519 + 85 + 36 + 16 = 656',\n",
       " 'correct': True,\n",
       " 'correct_reason': 'Solution is correct'}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test reward function\n",
    "from datasets import load_from_disk\n",
    "eval_dataset = load_from_disk(\"output/test_vllm__qwen2.5-1.5b-sft-distill-lora_results\")\n",
    "idx = 11\n",
    "eval_dataset[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Soft format reward: [0.5]\n",
      "Correctness reward: [2]\n"
     ]
    }
   ],
   "source": [
    "print(f'Soft format reward: {soft_format_reward_func([[{\"content\": eval_dataset[idx][\"completion\"]}]])}')\n",
    "score = correctness_reward_func(\n",
    "    [[{\"content\": eval_dataset[idx][\"completion\"]}]],\n",
    "    [eval_dataset[idx][\"numbers\"]],\n",
    "    [eval_dataset[idx][\"target\"]])\n",
    "print(f'Correctness reward: {score}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[32m\u001b[41mERROR\u001b[0m Failed to detect the name of this notebook. You can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mswulling\u001b[0m to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.19.9"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/data/countdown/wandb/run-20250422_211550-czswwj3v</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/swulling/countdown-rl-simple-10k/runs/czswwj3v' target=\"_blank\">iconic-forest-8</a></strong> to <a href='https://wandb.ai/swulling/countdown-rl-simple-10k' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/swulling/countdown-rl-simple-10k' target=\"_blank\">https://wandb.ai/swulling/countdown-rl-simple-10k</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/swulling/countdown-rl-simple-10k/runs/czswwj3v' target=\"_blank\">https://wandb.ai/swulling/countdown-rl-simple-10k/runs/czswwj3v</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<button onClick=\"this.nextSibling.style.display='block';this.style.display='none';\">Display W&B run</button><iframe src='https://wandb.ai/swulling/countdown-rl-simple-10k/runs/czswwj3v?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x7ff362ce16f0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wandb\n",
    "wandb.init(project=\"countdown-rl-simple-10k\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import GRPOConfig, GRPOTrainer\n",
    "training_args = GRPOConfig(\n",
    "    use_vllm = True, # use vLLM for fast inference!\n",
    "    learning_rate = 5e-6,\n",
    "    adam_beta1 = 0.9,\n",
    "    adam_beta2 = 0.99,\n",
    "    weight_decay = 0.1,\n",
    "    warmup_ratio = 0.01,\n",
    "    lr_scheduler_type = \"linear\",\n",
    "    optim = \"adamw_8bit\",\n",
    "    bf16 = is_bfloat16_supported(),\n",
    "    fp16 = not is_bfloat16_supported(),\n",
    "    per_device_train_batch_size = 8,\n",
    "    gradient_accumulation_steps = 4, # Increase to 4 for smoother training\n",
    "    num_generations = 8, # Decrease if out of memory\n",
    "    vllm_max_model_len= max_seq_length,\n",
    "    max_prompt_length = 512,\n",
    "    max_completion_length = 3584,\n",
    "    temperature = 1.0, # set to 1.0 for more diverse responses\n",
    "    #num_train_epochs = 1, # Set to 1 for a full training run\n",
    "    max_steps = 5000,\n",
    "    save_steps = 100,\n",
    "    max_grad_norm = 0.1,\n",
    "    output_dir = \"output/rl4\",\n",
    "    beta=0.001,\n",
    "    report_to = \"wandb\", # Can use Weights & Biases\n",
    "    log_completions=True,\n",
    "    logging_steps=1,\n",
    "    # unsloth grpo eval bug: https://github.com/unslothai/unsloth/issues/2367\n",
    "    # do_eval=True,\n",
    "    # eval_strategy=\"steps\",\n",
    "    # eval_steps=1,\n",
    "    # per_device_eval_batch_size = 32,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1\n",
      "   \\\\   /|    Num examples = 10,000 | Num Epochs = 2 | Total steps = 5,000\n",
      "O^O/ \\_/ \\    Batch size per device = 8 | Gradient accumulation steps = 4\n",
      "\\        /    Data Parallel GPUs = 1 | Total batch size (8 x 4 x 1) = 32\n",
      " \"-____-\"     Trainable parameters = 73,859,072/5,000,000,000 (1.48% trained)\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unsloth: Will smartly offload gradients to save VRAM!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<string>:1: SyntaxWarning: 'int' object is not callable; perhaps you missed a comma?\n",
      "<string>:1: SyntaxWarning: 'int' object is not callable; perhaps you missed a comma?\n",
      "<string>:1: SyntaxWarning: 'int' object is not callable; perhaps you missed a comma?\n"
     ]
    }
   ],
   "source": [
    "trainer = GRPOTrainer(\n",
    "    model = model,\n",
    "    processing_class = tokenizer,\n",
    "    reward_funcs = [\n",
    "        soft_format_reward_func,\n",
    "        correctness_reward_func,\n",
    "    ],\n",
    "    args = training_args,\n",
    "    train_dataset = dataset,\n",
    "    # eval_dataset = eval_dataset,\n",
    ")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "v4: https://wandb.ai/swulling/countdown-rl-simple-10k?nw=nwuserswulling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"output/qwen2.5-1.5b-rl-v4-lora\")  # Local saving lora weights\n",
    "tokenizer.save_pretrained(\"output/qwen2.5-1.5b-rl-v4-lora\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "```bash\n",
    "vllm serve output/models/Qwen2.5-1.5B-Instruct --port 8100 --api-key NLUKKXIJDZ91rpg1z --enforce-eager  --max-model-len 4096 --enable-lora --max-lora-rank 64 --lora-modules qwen2.5-1.5b-rl-v4-lora=output/qwen2.5-1.5b-rl-v4-lora\n",
    "\n",
    "CURATOR_VIEWER=1 python eval.py --provider vllm --data_path data/test_simple.jsonl --model_name qwen2.5-1.5b-rl-v4-lora --temperature 0.01 --max_tokens 2048\n",
    "\n",
    "https://curator.bespokelabs.ai/datasets/fc75d8c833ba4d7984d925604337a9d5  \n",
    "\n",
    "Accuracy: 45/100 (45.00%)\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
